/**
 * Copyright (C) 2010 - 2013 Harry Glasgow
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package com.googlecode.jaden.ui.dialog.automatedtraining;

import com.googlecode.jaden.common.config.LayerConfig;
import com.googlecode.jaden.common.config.ProjectConfiguration;
import com.googlecode.jaden.common.enums.LayerType;
import com.googlecode.jaden.common.other.JadenException;
import com.googlecode.jaden.common.other.Utils;
import com.googlecode.jaden.engine.*;
import com.googlecode.jaden.ui.Constants;
import com.googlecode.jaden.ui.UI;
import com.googlecode.jaden.ui.UiUtils;
import com.googlecode.jaden.ui.dialog.BaseDialog;

import javax.swing.*;
import java.awt.*;
import java.util.HashMap;
import java.util.Map;
import java.util.TreeMap;
import java.util.concurrent.*;
import java.util.prefs.Preferences;

import static com.googlecode.jaden.ui.Constants.*;

public class AutomatedTrainingDialog extends BaseDialog {

    private final UI ui;
    private final JLabel layersField;
    private final JTextField maskField;
    private final JTextField layerConfigField;
    private final JLabel roundLabel;
    private final JLabel jogLabel;
    private final JLabel statusLabel;
    private final JProgressBar runningBar;
    private final JList<String> bestList;
    private final DefaultComboBoxModel<String> bestComboBoxModel;
    private final java.util.Map<Double, AutomatedConfig> bestConfig;

    private final ExecutorService background;

    public AutomatedTrainingDialog(UI ui, Preferences preferences, ProjectConfiguration projectConfiguration) {
        super(ui, Constants.DIALOG_AUTOMATED_TRAINING_TITLE);
        this.ui = ui;
        JTabbedPane tabbedPane = new JTabbedPane();

        // Current
        layersField = new JLabel();
        maskField = new JTextField(40);
        maskField.setEditable(false);
        layerConfigField = new JTextField(40);
        layerConfigField.setEditable(false);
        roundLabel = new JLabel("1");
        jogLabel = new JLabel("0");
        JPanel currentPanel = createCurrentPanel();

        // Best
        bestComboBoxModel = new DefaultComboBoxModel<String>();
        bestList = new JList<String>(bestComboBoxModel);
        bestConfig = new TreeMap<Double, AutomatedConfig>();
        JPanel bestPanel = createBestPanel();

        // Tabs
        tabbedPane.addTab(Utils.translate(Constants.LABEL_CURRENT_TEXT), currentPanel);
        tabbedPane.addTab(Utils.translate(Constants.LABEL_BEST_TEXT), bestPanel);

        // Status
        statusLabel = new JLabel(Utils.translate(Constants.LABEL_READY_TEXT));
        runningBar = new JProgressBar();
        runningBar.setIndeterminate(true);
        runningBar.setVisible(false);

        // Main
        setMainPanel(UiUtils.verticalList(tabbedPane), preferences);
        setSouth(UiUtils.horizontalList(statusLabel, runningBar));

        // Process
        background = Executors.newSingleThreadExecutor();
        background.execute(new Gopher(projectConfiguration));

    }

    private JPanel createCurrentPanel() {
        JLabel inputMaskTextLabel = new JLabel(Utils.translate(Constants.LABEL_INPUT_MASK_TEXT));
        JLabel layersTextLabel = new JLabel(Utils.translate(Constants.LABEL_LAYERS_TEXT));
        JLabel configTextLabel = new JLabel(Utils.translate(Constants.LABEL_LAYER_CONFIG_TEXT));
        JLabel roundTextLabel = new JLabel(Utils.translate(Constants.LABEL_ROUND_TEXT));
        JLabel jogTextLabel = new JLabel(Utils.translate(Constants.LABEL_JOG_TEXT));
        UiUtils.enWidth(inputMaskTextLabel, layersTextLabel, configTextLabel, roundTextLabel, jogTextLabel);
        return UiUtils.verticalList(
                UiUtils.horizontalList(inputMaskTextLabel, maskField),
                UiUtils.horizontalList(layersTextLabel, layersField),
                UiUtils.horizontalList(configTextLabel, layerConfigField),
                UiUtils.horizontalList(roundTextLabel, roundLabel),
                UiUtils.horizontalList(jogTextLabel, jogLabel)
        );
    }


    private JPanel createBestPanel() {
        JScrollPane scroll = new JScrollPane(bestList);
        scroll.setPreferredSize(new Dimension(200, 125));
        return UiUtils.verticalList(
                UiUtils.horizontalList(scroll)
        );
    }

    protected void doOk() throws JadenException {
        if (saveBest()) {
            dispose();
        }
    }

    private boolean saveBest() {
        if (!bestConfig.isEmpty()) {
            int result = JOptionPane.showOptionDialog(this, Utils.translate(Constants.MESSAGE_SAVE_BEST_NETWORK_CONFIG),
                    Utils.translate(LABEL_UI_TITLE), JOptionPane.YES_NO_CANCEL_OPTION,
                    JOptionPane.QUESTION_MESSAGE, null, new Object[]{
                    Utils.translate(LABEL_YES_TEXT),
                    Utils.translate(LABEL_NO_TEXT),
                    Utils.translate(LABEL_CANCEL_TEXT)}, 0);
            if (result == 2) {
                return false; // Cancel
            } else if (result == 0) { // Yes
                // Save best.
                AutomatedConfig automatedConfig =
                        bestConfig.values().toArray(new AutomatedConfig[bestConfig.size()])[0];
                try {
                    ui.updateAutomatedTraining(automatedConfig.getMask(), automatedConfig.getLayerMap());
                    return true;
                } catch (Exception e) {
                    Utils.handleException(e);
                    return false;
                }
            }
        }
        return true;
    }

    @Override
    public void dispose() {
        try {
            background.shutdown();
        } catch (Exception e) {
            // Don't care.
        }
        super.dispose();
    }

    /**
     * =======================================
     * =======================================
     * Gopher - does all the background stuff.
     * =======================================
     * =======================================
     */
    private class Gopher implements Runnable {

        private final ProjectConfiguration projectConfiguration;

        private double bestError = Double.MAX_VALUE;
        private DataSet dataSet;

        private Gopher(ProjectConfiguration projectConfiguration) {
            this.projectConfiguration = projectConfiguration;
        }

        public void run() {
            try {
                runningBar.setVisible(true);
                statusLabel.setText(Utils.translate(Constants.LABEL_RUNNING_TEXT));

                dataSet = new DataSet(projectConfiguration);

                TrainingMode trainingMode = projectConfiguration.getTrainingMode();
                int numberOfInputs = projectConfiguration.getInputConfigurations().length;
                boolean[] mask = new boolean[numberOfInputs];
                if (trainingMode == TrainingMode.Forward) {
                    cycleForward(mask);
                } else if (trainingMode == TrainingMode.Backward) {
                    for (int i = 0; i < mask.length; i++) {
                        mask[i] = true;
                    }
                    cycleNumberOfLayers(mask);
                    cycleBackward(mask);
                } else if (trainingMode == TrainingMode.Exhaustive) {
                    cycleExhaustive(mask);
                } else if (trainingMode == TrainingMode.None) {
                    for (int i = 0; i < mask.length; i++) {
                        mask[i] = true;
                    }
                    cycleNumberOfLayers(mask);
                }
            } catch (JadenException e) {
                dispose();
                Utils.handleException(e);
                JOptionPane.showMessageDialog(AutomatedTrainingDialog.this, e.getMessage());
            } catch (Exception e) {
                dispose();
                Utils.handleException(e);
                JOptionPane.showMessageDialog(AutomatedTrainingDialog.this, Utils.translate(ERROR_GENERAL));
            } finally {
                runningBar.setVisible(false);
                statusLabel.setText(Utils.translate(Constants.LABEL_COMPLETE_TEXT));
            }
        }

        private void cycleForward(boolean[] mask) throws JadenException {
            int bestIndex = -1;
            for (int i = 0; i < mask.length; i++) {
                if (!mask[i]) {
                    double currentError = bestError;
                    mask[i] = true;
                    cycleNumberOfLayers(mask);
                    mask[i] = false;
                    if (bestError < currentError) {
                        bestIndex = i;
                    }
                }
            }
            if (bestIndex < 0) {
                return;
            }

            mask[bestIndex] = true;
            cycleForward(mask);
        }

        private void cycleBackward(boolean[] mask) throws JadenException {
            int bestIndex = -1;
            for (int i = 0; i < mask.length; i++) {
                if (mask[i]) {
                    double currentError = bestError;
                    mask[i] = false;
                    cycleNumberOfLayers(mask);
                    mask[i] = true;
                    if (bestError < currentError) {
                        bestIndex = i;
                    }
                }
            }
            if (bestIndex < 0) {
                return;
            }

            mask[bestIndex] = false;
            cycleBackward(mask);
        }

        private void cycleExhaustive(boolean[] mask) throws JadenException {
            long power = (long) Math.pow(2.0, mask.length);
            for (long i = 1; i < power; i++) {
                long shield = 1;
                for (int j = 0; j < mask.length; j++) {
                    mask[j] = (i & shield) != 0;
                    shield *= 2;
                }
                cycleNumberOfLayers(mask);
            }
        }

        /**
         * Adjust number of layers from min to max.
         */
        private void cycleNumberOfLayers(boolean[] mask) throws JadenException {

            // Set mask in UI.
            maskField.setText(displayMaskAsText(mask));

            for (int numberOfLayers = projectConfiguration.getMinimumLayers();
                 numberOfLayers <= projectConfiguration.getMaximumLayers();
                 numberOfLayers++) {

                layersField.setText(String.valueOf(numberOfLayers));
                cycleLayersWidth(mask, numberOfLayers);
            }
        }

        private String displayMaskAsText(boolean[] mask) {
            StringBuilder sb = new StringBuilder();
            for (boolean b : mask) {
                sb.append(b ? "1" : "0");
            }
            return sb.toString();
        }

        /**
         * Adjust the central layer width from 1 to (2 * active inputs width),
         * and adjust others accordingly.
         */
        private void cycleLayersWidth(boolean[] mask, int numberOfLayers) throws JadenException {
            int activeInputsWidth = 0;
            int inputPointer = 0;
            for (boolean b : mask) {
                activeInputsWidth += b ?
                        projectConfiguration.getInputConfigurations()[inputPointer].getInternalWidth() :
                        0;
                inputPointer++;
            }

            int outputWidth = projectConfiguration.getOutputConfiguration().getInternalWidth();

            if (numberOfLayers == 1) {
                // Only one option here, a singleton neuron.
                setupLayerWidths(mask, numberOfLayers, 0, activeInputsWidth, outputWidth);
            } else {
                for (int middleLayerWidth = 1;
                     middleLayerWidth <= 2 * (activeInputsWidth + outputWidth);
                     middleLayerWidth++) {
                    setupLayerWidths(mask, numberOfLayers, middleLayerWidth, activeInputsWidth, outputWidth);
                }
            }
        }

        private void setupLayerWidths(boolean[] mask, int numberOfLayers, int middleLayerWidth, int activeInputsWidth,
                                      int outputWidth) throws JadenException {
            int[] layerWidths;
            LayerType[] layerTypes;
            switch (numberOfLayers) {
                case 1:
                    layerWidths = new int[1];
                    layerWidths[0] = outputWidth;
                    layerTypes = new LayerType[1];
                    break;
                case 2:
                    layerWidths = new int[2];
                    layerWidths[0] = middleLayerWidth;
                    layerWidths[1] = outputWidth;
                    layerTypes = new LayerType[2];
                    break;
                case 3:
                    layerWidths = new int[3];
                    layerWidths[0] = (int) ((middleLayerWidth + activeInputsWidth) / 2.0); // Mid of middleLayer and inputs.
                    layerWidths[1] = middleLayerWidth;
                    layerWidths[2] = outputWidth;
                    layerTypes = new LayerType[3];
                    break;
                case 4:
                    layerWidths = new int[4];
                    layerWidths[0] = (int) ((middleLayerWidth + activeInputsWidth) / 2.0); // Mid of middleLayer and inputs.
                    layerWidths[1] = middleLayerWidth;
                    layerWidths[2] = (int) ((middleLayerWidth + outputWidth) / 2.0); // Mid of middleLayer and output.
                    layerWidths[3] = outputWidth;
                    layerTypes = new LayerType[4];
                    break;
                case 5:
                default:
                    layerWidths = new int[5];
                    layerWidths[0] = (int) ((middleLayerWidth + activeInputsWidth * 2.0) / 3.0); // Mid of middleLayer and inputs.
                    layerWidths[1] = (int) ((middleLayerWidth * 2.0 + activeInputsWidth) / 3.0); // Mid of middleLayer and inputs.
                    layerWidths[2] = middleLayerWidth;
                    layerWidths[3] = (int) ((middleLayerWidth + outputWidth) / 2.0); // Mid of middleLayer and output.
                    layerWidths[4] = outputWidth;
                    layerTypes = new LayerType[5];
                    break;
            }
            for (int i = 0; i < layerTypes.length; i++) {
                layerTypes[i] = LayerType.Soft;
            }
            cycleLayerTypes(mask, layerWidths, layerTypes);
        }

        private String displayLayersAsText(int[] layerWidths, LayerType[] layerTypes) {
            StringBuilder sb = new StringBuilder();
            for (int i = 0; i < layerWidths.length; i++) {
                sb.append(layerWidths[i]);
                sb.append(" ");
                sb.append(layerTypes[i]);
                if (i < layerWidths.length - 1) {
                    sb.append(" : ");
                }
            }
            return sb.toString();
        }

        private void cycleLayerTypes(boolean[] mask, int[] layerWidths, LayerType[] layerTypes) throws JadenException {

            final double pow = Math.pow(3, layerTypes.length);
            int layerTypePointer = 0;

            while (layerTypePointer < pow) {
                int layerTypePointerWorkingCopy = layerTypePointer;
                for (int index = 0; index < layerTypes.length; index++) {
                    int remainderThree = layerTypePointerWorkingCopy % 3;
                    switch (remainderThree) {
                        case 0:
                            layerTypes[index] = LayerType.Soft;
                            break;
                        case 1:
                            layerTypes[index] = LayerType.Normal;
                            break;
                        case 2:
                        default:
                            layerTypes[index] = LayerType.Sharp;
                            break;
                    }
                    layerTypePointerWorkingCopy /= 3;
                }

                layerConfigField.setText(displayLayersAsText(layerWidths, layerTypes));
                buildNetworkAndModel(mask, layerWidths, layerTypes);
                layerTypePointer++;
            }
        }

        private void buildNetworkAndModel(boolean[] mask, int[] layerWidths, LayerType[] layerTypes)
                throws JadenException {

            // Build an internal network based on inputsWidth and layerTypes.
            Map<Integer, Layer> layerMap = new HashMap<Integer, Layer>();
            int precedingSize = projectConfiguration.getInternalInputWidth();
            for (int i = 0; i < layerTypes.length; i++) {
                LayerConfig config = new LayerConfig(layerTypes[i], layerWidths[i]);
                layerMap.put(i, new Layer(config, precedingSize));
                precedingSize = config.getLayerWidth();
            }
            Network network = new Network(layerMap, projectConfiguration.getAutomatedLearningRate(),
                    projectConfiguration.getAutomatedMomentum());
            network.reset();

            // Set up IOConfigurations for the active inputs and the output.
            Model model = new Model(projectConfiguration.getInputConfigurations(), mask, network,
                    projectConfiguration.getOutputConfiguration());
            doRounds(mask, layerMap, model);
        }

        private void doRounds(boolean [] mask, Map<Integer, Layer> layerMap, Model model) throws JadenException {
            for (int round = 1; round <= projectConfiguration.getRounds(); round++) {
                roundLabel.setText(String.valueOf(round));
                int jog = 0;
                int jogs = projectConfiguration.getJogs();
                do {
                    jogLabel.setText(String.valueOf(jog));
                    doEpochs(mask, layerMap, model);
                    if (jog < jogs) {
                        model.jog();
                    }
                } while (jog++ < jogs);
            }
        }

        private void doEpochs(boolean [] mask, Map<Integer, Layer> layerMap, Model model) throws JadenException {
            for (int i = 0; i < projectConfiguration.getAutomatedEpochs(); i++) {
                doPass(mask, layerMap, model);
            }
        }

        private void doPass(boolean [] mask, Map<Integer, Layer> layerMap, Model model) throws JadenException {
            Trainer trainer = new Trainer(dataSet, model);
            trainer.train();
            double error = trainer.validate();
            if (error < bestError) {
                bestError = error;
                statusLabel.setText(Utils.translate(Constants.LABEL_RUNNING_BEST_ERROR_TEXT) + ' ' + bestError);
                AutomatedConfig automatedConfig = new AutomatedConfig(mask, layerMap);
                bestConfig.put(bestError, automatedConfig);
                bestComboBoxModel.insertElementAt(automatedConfigPlusErrorAsString(automatedConfig, bestError), 0);
            }
        }

        private String automatedConfigPlusErrorAsString(AutomatedConfig automatedConfig, double bestError) {
            return displayMaskAsText(automatedConfig.getMask()) + "/" + automatedConfig.getLayerMap().size() + " - " +
                    bestError;
        }
    }

    private class AutomatedConfig {

        private final boolean[] mask;
        private final Map<Integer, Layer> layerMap;

        private AutomatedConfig(boolean[] mask, Map<Integer, Layer> layerMap) {
            this.mask = mask;
            this.layerMap = layerMap;
        }

        public boolean[] getMask() {
            return mask;
        }

        public Map<Integer, Layer> getLayerMap() {
            return layerMap;
        }
    }
}
